package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"proxyman/utils"
	"sync"
)

// 定义代理服务器的配置
const (
	LocalAddr  = "127.0.0.1:6891"   // 本地监听地址
	RemoteAddr = "10.225.9.94:6891" // 目标服务器地址
)

func main() {
	// 开始监听本地地址
	listener, err := net.Listen("tcp", LocalAddr)
	if err != nil {
		log.Fatalf("Failed to listen on %s: %v", LocalAddr, err)
	}
	defer listener.Close()

	log.Printf("TCP reverse proxy server is running on %s, forwarding to %s", LocalAddr, RemoteAddr)

	for {
		// 接受客户端连接
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %v", err)
			continue
		}

		// 处理每个客户端连接
		go handleClient(clientConn, RemoteAddr)
	}
}

func handleClient(clientConn net.Conn, remoteAddr string) {
	defer clientConn.Close()

	// 连接到目标服务器
	remoteConn, err := net.Dial("tcp", remoteAddr)
	if err != nil {
		log.Printf("Failed to connect to remote server %s: %v", remoteAddr, err)
		return
	}
	defer remoteConn.Close()

	// 使用 WaitGroup 等待两个方向的数据传输完成
	var wg sync.WaitGroup
	wg.Add(2)

	// 定义一个通道，用于通知连接关闭
	closeChan := make(chan struct{})

	// 将客户端数据转发到目标服务器
	go func() {
		defer wg.Done()
		defer close(closeChan) // 通知关闭
		transferData(clientConn, remoteConn, &wg, ">>>")
	}()

	// 将目标服务器数据转发回客户端
	go func() {
		defer wg.Done()
		select {
		case <-closeChan:
			// 如果另一个方向已经关闭，则直接退出
			return
		default:
			transferData(remoteConn, clientConn, &wg, "<<<")
		}
	}()

	// 等待两个方向的数据传输完成
	wg.Wait()
}

// 数据传输函数
// func transferData(src, dst net.Conn, wg *sync.WaitGroup) {
// 	defer wg.Done()
// 	defer src.Close()

// 	_, err := io.Copy(dst, src)
// 	if err != nil {
// 		log.Printf("Error during data transfer: %v", err)
// 	}
// }

// 数据传输函数
func transferData(src, dst net.Conn, wg *sync.WaitGroup, direction string) {
	// defer wg.Done()
	defer src.Close()

	// 自定义 Writer，用于在写入时打印数据
	printingWriter := &PrintingWriter{
		Writer:    dst, // 将数据写入目标连接
		direction: direction,
	}

	// buf := new(bytes.Buffer)            // 用于存储读取的数据
	// teeReader := io.TeeReader(src, buf) // 同时将数据写入缓冲区

	written, err := io.Copy(printingWriter, src)
	if err != nil {
		log.Printf("Error during data transfer (%s): %v\n", direction, err)
	}

	// 打印传输的数据到日志
	// fmt.Printf("Transferred data (%s): %q\n", direction, buf.String())
	fmt.Printf("Transferred data (%s): %d bytes written\n", direction, written)

}

// 自定义 Writer，用于在写入时打印数据
type PrintingWriter struct {
	io.Writer
	direction string
}

// 实现 Write 方法，在写入时打印数据
func (pw *PrintingWriter) Write(p []byte) (n int, err error) {
	// 打印数据到终端
	// fmt.Printf("raw %s: %q\n", pw.direction, p)
	fmt.Printf("raw %s start ===\n", pw.direction)

	if len(p) > 2 {
		var dp []byte
		var dpk_dp string
		var err error

		if utils.IsZip(p) {
			dp, err = utils.DecompressZip(p)
		} else if utils.IsZip(p[2:]) {
			dp, err = utils.DecompressZip(p[2:])
		}

		if err == nil && dp != nil {
			dpk_dp, err = utils.GBKDecode(dp)
		}

		if err == nil && len(dpk_dp) > 0 {
			fmt.Printf("1 %s: %s\n", pw.direction, string(dpk_dp))
		} else {
			// 如果解压或解码失败，进入默认分支
			fmt.Printf("3 %s: %q\n", pw.direction, p)
		}

	} else {
		// length <= 2 的情况
		fmt.Printf("4 %s: %q\n", pw.direction, p)
	}

	fmt.Printf("raw %s end ===\n", pw.direction)
	// 将数据写入原始目标
	return pw.Writer.Write(p)
}
